package distributed

import (
	"context"
	"fmt"
	"sync"
	"time"

	proto "xagent/proto/v1"

	"github.com/asynkron/protoactor-go/actor"
)

// AdaptationManager handles distributed adaptation processes
type AdaptationManager struct {
	system         *actor.ActorSystem
	clusterManager *ClusterManager
	adaptations    map[string]*proto.DistributedAdaptation
	mutex          sync.RWMutex
}

// NewAdaptationManager creates a new adaptation manager
func NewAdaptationManager(system *actor.ActorSystem, clusterManager *ClusterManager) *AdaptationManager {
	return &AdaptationManager{
		system:         system,
		clusterManager: clusterManager,
		adaptations:    make(map[string]*proto.DistributedAdaptation),
	}
}

// InitiateAdaptation starts a new adaptation process
func (am *AdaptationManager) InitiateAdaptation(ctx context.Context, adaptationType string, currentState string) (*proto.DistributedAdaptation, error) {
	adaptation := &proto.DistributedAdaptation{
		Id:             fmt.Sprintf("adapt-%d", time.Now().UnixNano()),
		AdaptationType: adaptationType,
		CurrentState:   currentState,
		Parameters:     make(map[string]float32),
		Results:        make(map[string]float32),
		Status:         proto.AdaptationStatus_ADAPTATION_STATUS_EVALUATING,
		SuccessRate:    0,
		Timestamp:      time.Now().UnixNano(),
	}

	// Store adaptation
	am.mutex.Lock()
	am.adaptations[adaptation.Id] = adaptation
	am.mutex.Unlock()

	// Broadcast adaptation request
	if err := am.clusterManager.BroadcastMessage(adaptation); err != nil {
		return nil, fmt.Errorf("failed to broadcast adaptation: %v", err)
	}

	return adaptation, nil
}

// UpdateAdaptation updates an adaptation process
func (am *AdaptationManager) UpdateAdaptation(adaptationID string, results map[string]float32, status proto.AdaptationStatus) error {
	am.mutex.Lock()
	defer am.mutex.Unlock()

	adaptation, exists := am.adaptations[adaptationID]
	if !exists {
		return fmt.Errorf("adaptation %s not found", adaptationID)
	}

	// Update results and status
	for k, v := range results {
		adaptation.Results[k] = v
	}
	adaptation.Status = status

	// Calculate success rate if completed
	if status == proto.AdaptationStatus_ADAPTATION_STATUS_SUCCEEDED {
		adaptation.SuccessRate = calculateSuccessRate(adaptation.Results)
	}

	return nil
}

// GetAdaptation returns the current status of an adaptation
func (am *AdaptationManager) GetAdaptation(adaptationID string) *proto.DistributedAdaptation {
	am.mutex.RLock()
	defer am.mutex.RUnlock()

	return am.adaptations[adaptationID]
}

// ListAdaptations returns all adaptations
func (am *AdaptationManager) ListAdaptations() []*proto.DistributedAdaptation {
	am.mutex.RLock()
	defer am.mutex.RUnlock()

	adaptations := make([]*proto.DistributedAdaptation, 0, len(am.adaptations))
	for _, adaptation := range am.adaptations {
		adaptations = append(adaptations, adaptation)
	}
	return adaptations
}

// CleanupAdaptations removes completed or failed adaptations
func (am *AdaptationManager) CleanupAdaptations() {
	am.mutex.Lock()
	defer am.mutex.Unlock()

	for id, adaptation := range am.adaptations {
		if adaptation.Status == proto.AdaptationStatus_ADAPTATION_STATUS_SUCCEEDED ||
			adaptation.Status == proto.AdaptationStatus_ADAPTATION_STATUS_FAILED {
			delete(am.adaptations, id)
		}
	}
}

// calculateSuccessRate calculates the success rate from adaptation results
func calculateSuccessRate(results map[string]float32) float32 {
	if len(results) == 0 {
		return 0
	}

	var totalScore float32
	for _, score := range results {
		totalScore += score
	}

	return totalScore / float32(len(results))
}
